import numpy as np
from scipy import linalg
import control

class LQRController:
    """
    LQR控制器 - 离散时间线性二次型调节器
    """
    def __init__(self, A, B, Q=None, R=None):
        """
        初始化LQR控制器
        A: 系统状态矩阵
        B: 控制输入矩阵
        Q: 状态权重矩阵，默认为单位阵
        R: 控制权重矩阵，默认为单位阵
        """
        # 确保所有矩阵为浮点数类型
        self.A = A.astype(np.float64)
        self.B = B.astype(np.float64)
        
        # 检查系统的可控性
        if not self._is_controllable():
            print("警告：系统可能不完全可控，LQR性能可能不佳")
        
        # 状态权重矩阵 Q (正定)
        if Q is None:
            # 默认使用高位置权重和中等速度权重
            self.Q = np.diag([1000.0, 1000.0, 100.0, 100.0])
        else:
            self.Q = Q.astype(np.float64)
            
        # 确保Q是对称正定的
        self.Q = (self.Q + self.Q.T) / 2  # 确保对称
        
        # 控制权重矩阵 R (正定)
        if R is None:
            # 默认使用小的控制权重
            self.R = np.diag([0.01, 0.01])
        else:
            self.R = R.astype(np.float64)
            
        # 确保R是对称正定的
        self.R = (self.R + self.R.T) / 2  # 确保对称
        
        # 验证Q和R的正定性
        self._validate_matrices()
        
        # 计算LQR增益矩阵
        self.K = self._compute_lqr_gain()
        
        # 打印控制器信息
        print("LQR控制器初始化:")
        print("状态权重矩阵 Q:")
        print(self.Q)
        print("控制权重矩阵 R:")
        print(self.R)
        print("LQR增益矩阵 K:")
        print(self.K)
        
    def _validate_matrices(self):
        """验证Q和R矩阵的正定性"""
        try:
            # 尝试对Q进行Cholesky分解
            linalg.cholesky(self.Q)
        except linalg.LinAlgError:
            print("警告：Q矩阵不是正定的，已调整为正定矩阵")
            # 添加少量对角元素确保正定性
            self.Q = self.Q + np.eye(self.Q.shape[0]) * 1e-5
            
        try:
            # 尝试对R进行Cholesky分解
            linalg.cholesky(self.R)
        except linalg.LinAlgError:
            print("警告：R矩阵不是正定的，已调整为正定矩阵")
            # 添加少量对角元素确保正定性
            self.R = self.R + np.eye(self.R.shape[0]) * 1e-5
            
    def _is_controllable(self):
        """检查系统是否可控"""
        n = self.A.shape[0]  # 状态维数
        
        # 构建可控性矩阵
        C = np.zeros((n, n*self.B.shape[1]))
        temp = self.B.copy()
        C[:, :self.B.shape[1]] = temp
        
        for i in range(1, n):
            temp = np.dot(self.A, temp)
            C[:, i*self.B.shape[1]:(i+1)*self.B.shape[1]] = temp
            
        # 检查秩
        rank = np.linalg.matrix_rank(C)
        return rank == n
        
    def _compute_lqr_gain(self):
        """计算LQR增益矩阵K - 使用多种方法确保计算成功"""
        # 方法1：使用scipy的离散Riccati方程求解器
        try:
            P = linalg.solve_discrete_are(self.A, self.B, self.Q, self.R)
            # 计算增益矩阵K
            K = np.dot(np.dot(np.linalg.inv(self.R + np.dot(np.dot(self.B.T, P), self.B)), 
                              np.dot(self.B.T, P)), self.A)
            print("成功使用scipy的离散Riccati方程求解器")
            return K
        except Exception as e:
            print(f"scipy离散Riccati方程求解失败: {e}")
            
        # 方法2：使用control库的dlqr函数
        try:
            K, _, _ = control.dlqr(self.A, self.B, self.Q, self.R)
            print("成功使用control库的dlqr函数")
            return K
        except Exception as e:
            print(f"control.dlqr求解失败: {e}")
            
        # 方法3：使用迭代方法求解Riccati方程
        print("尝试使用迭代方法求解Riccati方程...")
        P = self._solve_riccati_iterative()
        K = np.dot(np.dot(np.linalg.inv(self.R + np.dot(np.dot(self.B.T, P), self.B)), 
                          np.dot(self.B.T, P)), self.A)
        return K
        
    def _solve_riccati_iterative(self, max_iter=1000, tol=1e-8):
        """迭代求解离散时间Riccati方程"""
        n = self.A.shape[0]
        P = np.eye(n)  # 初始化P为单位矩阵
        
        for i in range(max_iter):
            # 离散时间Riccati方程迭代
            K = np.dot(np.dot(np.linalg.inv(self.R + np.dot(np.dot(self.B.T, P), self.B)), 
                              np.dot(self.B.T, P)), self.A)
                              
            A_cl = self.A - np.dot(self.B, K)
            P_new = self.Q + np.dot(K.T, np.dot(self.R, K)) + np.dot(A_cl.T, np.dot(P, A_cl))
            
            # 检查收敛
            if np.linalg.norm(P_new - P) < tol:
                print(f"Riccati方程迭代收敛，迭代次数: {i+1}")
                return P_new
                
            P = P_new
            
        print(f"警告：Riccati方程迭代未收敛，使用最后一次迭代结果")
        return P
    
    def compute_control_input(self, current_state, reference_state):
        """
        计算控制输入
        current_state: 当前状态 [x, y, vx, vy]
        reference_state: 参考状态 [x_ref, y_ref, vx_ref, vy_ref]
        返回: 控制输入 [ax, ay]
        """
        # 确保状态为浮点数
        current_state = np.array(current_state, dtype=np.float64)
        reference_state = np.array(reference_state, dtype=np.float64)
        
        # 计算状态误差 (修正：当前状态 - 参考状态)
        state_error = current_state - reference_state
        
        # 打印当前状态与参考状态
        if np.linalg.norm(state_error) > 100:
            print(f"警告：状态误差过大 {np.linalg.norm(state_error):.2f}")
            print(f"当前状态: {current_state}")
            print(f"参考状态: {reference_state}")
        
        # 计算控制输入 u = -K * state_error
        u = -np.dot(self.K, state_error)
        
        # 限制控制输入（防止过大的控制输入）
        u_clipped = np.clip(u, -20.0, 20.0)
        if not np.array_equal(u, u_clipped):
            print(f"控制输入被限制: {u} -> {u_clipped}")
        
        return u_clipped
    
    def update_parameters(self, A=None, B=None, Q=None, R=None):
        """
        更新控制器参数
        A: 新的系统状态矩阵（可选）
        B: 新的控制输入矩阵（可选）
        Q: 新的状态权重矩阵（可选）
        R: 新的控制权重矩阵（可选）
        """
        if A is not None:
            self.A = A.astype(np.float64)
        if B is not None:
            self.B = B.astype(np.float64)
        if Q is not None:
            self.Q = Q.astype(np.float64)
            self.Q = (self.Q + self.Q.T) / 2  # 确保对称
        if R is not None:
            self.R = R.astype(np.float64)
            self.R = (self.R + self.R.T) / 2  # 确保对称
            
        # 验证矩阵
        self._validate_matrices()
        
        # 重新计算LQR增益矩阵
        self.K = self._compute_lqr_gain()
        
        print("LQR控制器参数已更新:")
        print("新的LQR增益矩阵 K:")
        print(self.K) 